import math
from time import sleep
import numpy as np
import matplotlib.pyplot as plt

stage1_list =  np.array([])
stage2_list =  np.array([])

def exp(yi, d, f):
    global stage1_list
    global stage2_list
    if abs(f - 0) > 1e-6:
        yi = yi - f

    # 计算yi * log2e
    A = yi
    # B = A >> 1
    B = A * math.pow(2, -1) 
    # D = yi >> 4
    D = A * math.pow(2, -4)
    E = A + B - D
    # print("E: ", E)
    # print("Rel:", yi * math.log(math.e, 2))

    # 获取ui， vi， ui为整数， vi 为浮点数
    vi, ui  = math.modf(E)
    
    stage1_list = np.append(stage1_list, (vi))
    stage2_list = np.append(stage2_list, math.pow(2, (vi)))

    # print("vi: ", vi)
    # 计算pow(2, vi)    
    # vi = abs(vi)
    f_vi = vi + d
    exp_yi = f_vi * math.pow(2, ui)
    
    # if ui > 0:
    #     exp_yi = f_vi * math.pow(2, ui)
    # else:
    #     exp_yi = f_vi / math.pow(2, abs(ui))


    # 计算
    if abs(exp_yi - math.exp(yi)) > 1e-5:
        print("yi: ", yi)
        print("ui: ", ui)
        print("h_exp: ", exp_yi)
        print("r_exp: ", math.exp(yi))
    return exp_yi


def ln(f):
    org_f = f
    # Get (k, w)
    w = 0
    vi, ui = math.modf(f)

    # ui >=2
    if f >= 2:
        while f >=2:
            f = f / 2.0
            w = w + 1
            ki, _ = math.modf(f)
    
    elif f < 1.0:
        while f < 1.0:
            f = f * 2.0
            w = w-1 
            ki, _ = math.modf(f)
    else:
        ki, _ = math.modf(f)

    k = ki + 1
    # print("ln k: ", k , " ln w: ", w)

    z = k -1 + w
    # 乘 ln2 = 0.1011
    ret = z * math.pow(2, -1) + z * math.pow(2, -3) + z * math.pow(2, -4)
    
    print("Rel: ", math.log(org_f))
    print("lnh: ", ret)
    return ret

# x_array 为nump数组
def softmax(x_array, d1, d2):
    result_arr = np.empty(x_array.shape)
    #获取F
    F = 0
    # for x in np.nditer(x_array):
    for i in range(x_array.size):
        F = exp(x_array[i], d1, 0) + F

    lnF = ln(F)

    # 更新数组
    for i in range(x_array.size):
        result_arr[i] = exp(x_array[i], d2, lnF)

    return result_arr

def softmax_soft(x_array):
    F = 0.0
    result_arr = np.empty(x_array.shape)
    for i in range(x_array.size):
        F = math.exp(x_array[i]) + F

    # 更新数组
    for i in range(x_array.size):
        result_arr[i] = math.exp(x_array[i]) / F
    
    return result_arr


if __name__=='__main__':
    
    nums = 2048

    # # uint 8
    # x_array = np.random.randint(0, 256, size=nums, dtype='int')
    # h_result = softmax(x_array, 0.9423, 0.9445)
    # s_result = softmax_soft(x_array)
    # diff = abs(h_result - s_result) < 8e-5
    # acc = np.average(diff)    
    # print("uint: ", acc)

    # # int 8
    # x_array = np.random.randint(-127, 127, size=nums, dtype='int')
    # h_result = softmax(x_array, 0.9423, 0.9445)
    # s_result = softmax_soft(x_array)
    # diff = abs(h_result - s_result) < 8e-5
    # acc = np.average(diff)    
    # print("int8: ", acc)
    
    # x_array = np.random.randint(-10, 10, size=nums, dtype='int')
    # h_result = softmax(x_array, 0.9423, 0.9445)
    # s_result = softmax_soft(x_array)
    # diff = abs(h_result - s_result) < 8e-5
    # acc = np.average(diff)    
    # print("round10: ", acc)
    
    x_array = np.random.randint(-128, 127, size=nums, dtype='int')
    x_array = x_array - np.max(x_array)
    h_result = softmax(x_array, 1.05, 1.15)
    s_result = softmax_soft(x_array)
    diff = abs(h_result - s_result) < 8e-5
    acc = np.average(diff)    
    print("round5: ", acc)
    
    
    step1 = stage1_list[:1024]
    step2 = stage2_list[:1024]
    avg1 = np.average(step1)
    avg2 = np.average(step2)
    print("vi average: ", avg1)
    print("ex average: ", avg2)
    print("diff: ", avg2 - avg1)
    
    
    step1 = stage1_list[1024:]
    step2 = stage2_list[1024:]
    avg1 = np.average(step1)
    avg2 = np.average(step2)
    print("vi average: ", avg1)
    print("ex average: ", avg2)
    print("diff: ", avg2 - avg1)
    
    
    # 显示图形
    # plt.subplot(2,1,1)
    # x = np.arange(0, 4096)
    # plt.plot(x, stage1_list)
    # plt.title('vi')

    # plt.subplot(2,1,2)
    # plt.plot(x, stage2_list)
    # plt.title('ex')
    # plt.show()
    











